# -*- coding: utf-8 -*-
"""
车牌识别服务核心模块

提供基于YOLOv5和深度学习的车牌检测和识别功能，支持：
- 单层和双层车牌识别
- 车牌颜色识别  
- 多种输入格式（本地文件、URL、numpy数组）
- 可配置的识别参数
- GPU加速支持

Author: AI Assistant
Date: 2025-06-12
"""

import sys
from pathlib import Path
import numpy as np
import torch
import time
import cv2
from easydict import EasyDict as edict
import re
import os
from urllib.parse import quote
import requests

from LogHelper import logger

# 添加YOLOv5模块路径
FILE = Path(__file__).resolve()
PARENT_DIR = FILE.parents[0]
YOLO_MODULE_PATH = str(PARENT_DIR / 'yolov5_plate')
sys.path.insert(0, YOLO_MODULE_PATH)

# YOLOv5相关导入
from yolov5_plate.models.experimental import attempt_load
from yolov5_plate.utils.datasets import letterbox
from yolov5_plate.utils.general import check_img_size, non_max_suppression_face, scale_coords
from yolov5_plate.utils.torch_utils import select_device, time_synchronized
from yolov5_plate.utils.cv_puttext import cv2ImgAddText

# 车牌识别相关导入
from plate_recognition.plate_rec import get_plate_result, init_model
from plate_recognition.double_plate_split_merge import get_split_merge
from plate_recognition.color_rec import plate_color_rec, init_color_model

# 默认配置参数
DEFAULT_CONFIG = edict({
    "yolo_weights": YOLO_MODULE_PATH + "/best.pt",
    "save_vid": "",
    "save_txt": "",
    "output": "inference/output",
    "img_size": 640,
    "conf_thres": 0.25,
    "iou_thres": 0.45,
    "fourcc": "mp4v",
    "device": '1',   # GPU设备ID或'cpu'
    "classes": None,
    "agnostic_nms": True,  # 在多个类别中应用NMS
    "augment": "",
    "evaluate": False,
    "rec_model": str(PARENT_DIR/"plate_recognition/model/plate_rec.pth"),
    "color_model": str(PARENT_DIR/"plate_recognition/model/color_classify.pth")
})

# 车牌格式验证正则表达式
LICENSE_PLATE_PATTERNS = [
    r'^[\u4e00-\u9fa5][A-Z][0-9A-Z]{5}$',  # 蓝牌格式：7位字符
    r'^[\u4e00-\u9fa5][A-Z][0-9A-Z]{6}$',  # 绿牌格式：8位字符
]

def scale_landmark_coordinates(img1_shape, coords, img0_shape, ratio_pad=None):
    """
    将坐标从一个图像尺寸缩放到另一个图像尺寸
    
    Args:
        img1_shape: 源图像形状 (height, width)
        coords: 需要缩放的坐标数组 (N, 8) - 4个关键点的x,y坐标
        img0_shape: 目标图像形状 (height, width)
        ratio_pad: 可选的缩放比例和填充信息
    
    Returns:
        缩放后的坐标数组
    """
    if ratio_pad is None:
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]

    # 调整x坐标
    coords[:, [0, 2, 4, 6]] -= pad[0]
    # 调整y坐标
    coords[:, [1, 3, 5, 7]] -= pad[1]
    # 应用缩放
    coords[:, :8] /= gain
    
    # 限制坐标在图像边界内
    coords[:, 0].clamp_(0, img0_shape[1])  # x1
    coords[:, 1].clamp_(0, img0_shape[0])  # y1
    coords[:, 2].clamp_(0, img0_shape[1])  # x2
    coords[:, 3].clamp_(0, img0_shape[0])  # y2
    coords[:, 4].clamp_(0, img0_shape[1])  # x3
    coords[:, 5].clamp_(0, img0_shape[0])  # y3
    coords[:, 6].clamp_(0, img0_shape[1])  # x4
    coords[:, 7].clamp_(0, img0_shape[0])  # y4
    
    return coords

def apply_perspective_transform(image, corner_points):
    """
    对图像应用透视变换，矫正车牌区域为矩形
    
    Args:
        image: 输入图像
        corner_points: 车牌四个角点坐标
    
    Returns:
        透视变换后的车牌图像
    """
    # 排序角点：左上、右上、右下、左下
    ordered_points = sort_corner_points(corner_points)
    top_left, top_right, bottom_right, bottom_left = ordered_points

    # 计算变换后的宽度
    width_bottom = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + 
                          ((bottom_right[1] - bottom_left[1]) ** 2))
    width_top = np.sqrt(((top_right[0] - top_left[0]) ** 2) + 
                       ((top_right[1] - top_left[1]) ** 2))
    max_width = max(int(width_bottom), int(width_top))

    # 计算变换后的高度
    height_right = np.sqrt(((top_right[0] - bottom_right[0]) ** 2) + 
                          ((top_right[1] - bottom_right[1]) ** 2))
    height_left = np.sqrt(((top_left[0] - bottom_left[0]) ** 2) + 
                         ((top_left[1] - bottom_left[1]) ** 2))
    max_height = max(int(height_right), int(height_left))

    # 定义目标矩形角点
    destination_points = np.array([
        [0, 0],
        [max_width - 1, 0],
        [max_width - 1, max_height - 1],
        [0, max_height - 1]
    ], dtype="float32")

    # 计算透视变换矩阵并应用
    transform_matrix = cv2.getPerspectiveTransform(ordered_points, destination_points)
    warped_image = cv2.warpPerspective(image, transform_matrix, (max_width, max_height))

    return warped_image

def sort_corner_points(points):
    """
    对四个角点进行排序：左上、右上、右下、左下
    
    Args:
        points: 4个点的坐标数组
    
    Returns:
        排序后的角点数组
    """
    rect = np.zeros((4, 2), dtype="float32")
    
    # 计算点的坐标和，最小和为左上角，最大和为右下角
    coord_sum = points.sum(axis=1)
    rect[0] = points[np.argmin(coord_sum)]    # 左上角
    rect[2] = points[np.argmax(coord_sum)]    # 右下角

    # 计算点的坐标差，最小差为右上角，最大差为左下角
    coord_diff = np.diff(points, axis=1)
    rect[1] = points[np.argmin(coord_diff)]   # 右上角
    rect[3] = points[np.argmax(coord_diff)]   # 左下角

    return rect

class LicensePlateConfig:
    """
    车牌识别配置类
    
    管理车牌识别过程中的各种参数设置，包括：
    - 图像保存选项
    - 检测和识别参数
    - 输出格式控制
    - 优先级设置等
    """
    
    def __init__(self, config_dict=None):
        """
        初始化配置参数
        
        Args:
            config_dict: 可选的配置字典，用于覆盖默认值
        """
        # 基础设置
        self.image_save = 0                    # 是否保存识别图片
        self.return_box = 0                    # 是否返回检测框坐标
        self.filter_error_plate = 1           # 是否过滤不符合规则的车牌
        self.return_src = 0                    # 是否返回原始识别结果
        
        # 优先级设置
        self.priority_province_chinese = None  # 优先省份汉字
        self.priority_province_char = None     # 优先省份字母
        self.chinese_thres = 0                 # 替换优先省份的阈值 (0-100)
        self.chinese_thres_for_char = 0        # 根据优先省份字母替换的阈值 (0-100)
        
        # 处理流程控制
        self.need_detect = 1                   # 是否需要进行检测
        self.need_transformer = 1              # 是否需要透视变换
        self.need_color = 1                    # 是否需要颜色识别
        
        # 模型设置
        self.detect_model = 'mtcnn'           # 检测模型类型
        self.plate_resize_type = 'resize'      # 车牌缩放类型
        self.lpr_type = 'lprnet'              # 识别方法类型
        
        # 显示设置
        self.show = False                      # 是否显示结果
        self.roi_box = None                   # 感兴趣区域

        # 应用用户自定义配置
        if config_dict:
            for key, value in config_dict.items():
                if hasattr(self, key):
                    setattr(self, key, value)

class LicensePlateRecognitionService:
    """
    车牌识别服务主类
    
    提供完整的车牌检测和识别功能，包括：
    - 模型初始化和管理
    - 图像预处理
    - 车牌检测
    - 车牌识别
    - 结果后处理
    """
    
    def __init__(self):
        """初始化车牌识别服务"""
        self._initialize_device()
        self._load_models()
        self._setup_parameters()
        logger.info("车牌识别服务初始化完成")

    def _initialize_device(self):
        """初始化计算设备（CPU或GPU）"""
        gpu_count = torch.cuda.device_count()
        if gpu_count > 0:
            device_id = str(gpu_count - 1)  # 使用最后一块GPU
            logger.info(f"检测到 {gpu_count} 块GPU，使用GPU {device_id}")
        else:
            device_id = 'cpu'
            logger.info("未检测到GPU，使用CPU")
            
        self.device = select_device(device_id)
        self.use_half_precision = self.device.type != 'cpu'

    def _load_models(self):
        """加载YOLOv5检测模型和车牌识别模型"""
        try:
            # 加载YOLOv5车牌检测模型
            self.detection_model = attempt_load(DEFAULT_CONFIG.yolo_weights, map_location=self.device)
            self.model_stride = int(self.detection_model.stride.max())
            self.image_size = check_img_size(DEFAULT_CONFIG.img_size, s=self.model_stride)
            self.class_names = (self.detection_model.module.names if hasattr(self.detection_model, 'module') 
                               else self.detection_model.names)
            
            if self.use_half_precision:
                self.detection_model.half()  # 转换为FP16以提高性能
                
            # 预热模型
            if self.device.type != 'cpu':
                dummy_input = torch.zeros(1, 3, self.image_size, self.image_size).to(self.device)
                dummy_input = dummy_input.type_as(next(self.detection_model.parameters()))
                self.detection_model(dummy_input)
                
            logger.info("YOLOv5检测模型加载成功")
            
            # 加载车牌识别模型
            self.plate_recognition_model = init_model(self.device, DEFAULT_CONFIG.rec_model)
            logger.info("车牌识别模型加载成功")
            
            # 加载车牌颜色识别模型
            self.plate_color_model = init_color_model(DEFAULT_CONFIG.color_model, self.device)
            logger.info("车牌颜色识别模型加载成功")
            
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            raise

    def _setup_parameters(self):
        """设置识别参数"""
        self.confidence_threshold = DEFAULT_CONFIG.conf_thres
        self.iou_threshold = DEFAULT_CONFIG.iou_thres
        self.target_classes = DEFAULT_CONFIG.classes
        self.agnostic_nms = DEFAULT_CONFIG.agnostic_nms

    def load_image_from_url(self, url, save_to_disk=False):
        """
        从URL加载图像
        
        Args:
            url: 图像URL
            save_to_disk: 是否保存到本地
            
        Returns:
            OpenCV格式的图像数据
        """
        try:
            # URL编码处理特殊字符
            encoded_url = quote(url, safe=":/?#[]@!$&'()*+,;=")
            
            response = requests.get(encoded_url, timeout=10)
            response.raise_for_status()
            
            # 将字节流转换为OpenCV图像
            image_array = np.frombuffer(response.content, dtype=np.uint8)
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            
            if image is None:
                raise ValueError("无法解码图像数据")
                
            if save_to_disk:
                self._save_image_to_disk(image, url)
                
            return image
            
        except requests.RequestException as e:
            raise Exception(f"从URL加载图像失败: {url}, 错误: {e}")
        except Exception as e:
            raise Exception(f"图像处理失败: {e}")

    def _save_image_to_disk(self, image, url):
        """将图像保存到磁盘"""
        try:
            filename = url.split("/")[-1] or "downloaded_image.jpg"
            save_dir = "./upload"
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, filename)
            cv2.imwrite(save_path, image)
            logger.info(f"图像已保存到: {save_path}")
        except Exception as e:
            logger.warning(f"保存图像失败: {e}")

    def recognize_license_plates(self, image_input, show_result=False, roi_box=None, config=None):
        """
        识别图像中的车牌
        
        Args:
            image_input: 图像输入，可以是numpy数组、文件路径或URL
            show_result: 是否显示识别结果
            roi_box: 感兴趣区域 [x1, y1, x2, y2]
            config: LicensePlateConfig配置对象
            
        Returns:
            车牌识别结果列表，每个结果包含车牌号、颜色、位置等信息
        """
        start_time = time.time()
        
        # 初始化配置
        if config is None:
            config = LicensePlateConfig()
        config.show = show_result
        config.roi_box = roi_box

        # 加载图像
        image = self._load_input_image(image_input, config)
        if image is None:
            logger.warning("图像加载失败")
            return []

        # 应用ROI裁剪
        if roi_box is not None:
            image = self._apply_roi_cropping(image, roi_box)

        load_time = time.time() - start_time
        logger.info(f"图像加载耗时: {load_time:.3f}秒")

        return self._perform_recognition(image, config)

    def _load_input_image(self, image_input, config):
        """加载不同类型的图像输入"""
        if isinstance(image_input, np.ndarray):
            return image_input
        elif isinstance(image_input, str):
            if image_input.startswith("http"):
                return self.load_image_from_url(image_input, save_to_disk=config.image_save)
            else:
                image = cv2.imread(image_input)
                if image is None:
                    logger.error(f"无法读取图像文件: {image_input}")
                return image
        else:
            logger.error(f"不支持的图像输入类型: {type(image_input)}")
            return None

    def _apply_roi_cropping(self, image, roi_box):
        """应用感兴趣区域裁剪"""
        try:
            x1, y1, x2, y2 = [max(0, int(coord)) for coord in roi_box]
            height, width = image.shape[:2]
            x2 = min(x2, width)
            y2 = min(y2, height)
            
            if x1 < x2 and y1 < y2:
                return image[y1:y2, x1:x2]
            else:
                logger.warning("ROI区域无效，使用原始图像")
                return image
        except Exception as e:
            logger.warning(f"ROI裁剪失败: {e}，使用原始图像")
            return image

    def _perform_recognition(self, image, config):
        """执行车牌识别"""
        if config.need_detect == 0:
            # 跳过检测，直接识别整个图像作为车牌
            height, width = image.shape[:2]
            detection_box = [0, 0, width, height, 1.0, 0, 0, width, 0, width, height, 0, height, 0]
            return self._recognize_plate_from_detection(image, detection_box, config)

        return self._detect_and_recognize(image, config)

    def _detect_and_recognize(self, image, config):
        """检测并识别车牌"""
        # 图像预处理
        processed_image = letterbox(image, new_shape=self.image_size)[0]
        processed_image = processed_image[:, :, ::-1].transpose(2, 0, 1)  # BGR转RGB，HWC转CHW
        processed_image = np.ascontiguousarray(processed_image)

        # 转换为PyTorch张量
        input_tensor = torch.from_numpy(processed_image).to(self.device)
        input_tensor = input_tensor.half() if self.use_half_precision else input_tensor.float()
        input_tensor /= 255.0  # 归一化到[0,1]
        
        if input_tensor.ndimension() == 3:
            input_tensor = input_tensor.unsqueeze(0)

        # 执行检测
        detection_start = time_synchronized()
        with torch.no_grad():
            detections = self.detection_model(input_tensor, augment=DEFAULT_CONFIG.augment)[0]

        # 应用NMS
        detections = non_max_suppression_face(
            detections, self.confidence_threshold, self.iou_threshold, 
            classes=self.target_classes, agnostic=self.agnostic_nms
        )
        
        detection_time = time_synchronized() - detection_start
        logger.info(f"车牌检测耗时: {detection_time:.3f}秒")

        # 处理检测结果
        recognition_start = time_synchronized()
        recognition_results = []
        
        for detection_batch in detections:
            if detection_batch is not None and len(detection_batch):
                # 坐标变换
                detection_batch[:, :4] = scale_coords(
                    input_tensor.shape[2:], detection_batch[:, :4], image.shape
                ).round()
                detection_batch[:, 5:13] = scale_landmark_coordinates(
                    input_tensor.shape[2:], detection_batch[:, 5:13], image.shape
                ).round()

                # 处理每个检测结果
                for detection in detection_batch.tolist():
                    detection[:4] = [int(x) for x in detection[:4]]
                    class_id = int(detection[13])
                    class_name = self.class_names[class_id]
                    
                    if class_name in ["single", "double"]:
                        plate_results = self._recognize_plate_from_detection(image, detection, config)
                        recognition_results.extend(plate_results)

                        # 可视化结果
                        if config.show:
                            self._visualize_results(image, plate_results)

        recognition_time = time_synchronized() - recognition_start
        logger.info(f"车牌识别耗时: {recognition_time:.3f}秒")

        if config.show:
            self._display_result_image(image)

        return recognition_results

    def _recognize_plate_from_detection(self, image, detection_box, config):
        """从检测框中识别车牌"""
        try:
            # 解析检测结果
            bbox = detection_box[:4]  # 边界框坐标
            confidence = detection_box[4]  # 置信度
            landmarks = detection_box[5:13]  # 关键点坐标
            plate_type = int(detection_box[13])  # 车牌类型

            # 构建关键点数组
            corner_points = np.zeros((4, 2))
            for i in range(4):
                corner_points[i] = [int(landmarks[2 * i]), int(landmarks[2 * i + 1])]

            # 透视变换获取车牌区域
            plate_roi = apply_perspective_transform(image, corner_points)
            
            # 处理双层车牌
            if plate_type == 1:
                plate_roi = get_split_merge(plate_roi)

            # 车牌字符识别
            recognition_start = time.time()
            plate_text, text_confidence, char_confidences = get_plate_result(
                plate_roi, self.device, self.plate_recognition_model
            )
            
            # 车牌颜色识别
            plate_color = plate_color_rec(plate_roi, self.plate_color_model, self.device)
            logger.info(f"识别车牌颜色: {plate_color}")

            # 应用优先省份规则
            plate_text, text_confidence, char_confidences = self._apply_province_priority(
                plate_text, text_confidence, char_confidences, config
            )

            # 验证车牌格式
            if not self._validate_plate_format(plate_text):
                logger.info(f"车牌格式无效: {plate_text}")
                return []

            # 调整颜色识别结果
            plate_color = self._adjust_plate_color(plate_text, plate_color)

            # 构建识别结果
            result = self._build_recognition_result(
                plate_text, plate_color, bbox, text_confidence, char_confidences
            )

            # 质量过滤
            if self._passes_quality_check(char_confidences):
                return [result]
            else:
                logger.info(f"车牌质量不满足要求: {plate_text}")
                return []

        except Exception as e:
            logger.error(f"车牌识别失败: {e}")
            return []

    def _apply_province_priority(self, plate_text, confidence, char_confidences, config):
        """应用优先省份规则"""
        if (config.priority_province_char and config.priority_province_chinese and 
            plate_text and plate_text[0] == config.priority_province_char):
            
            plate_text = config.priority_province_chinese + plate_text
            char_confidences.insert(0, 60)
            confidence *= 0.6
            
        return plate_text, confidence, char_confidences

    def _validate_plate_format(self, plate_text):
        """验证车牌格式是否符合规范"""
        if not plate_text:
            return False
            
        # 检查标准格式
        for pattern in LICENSE_PLATE_PATTERNS:
            if re.match(pattern, plate_text):
                return True
                
        # 检查特殊格式（学车、警车）
        if len(plate_text) == 7 and plate_text[-1] in ["学", "警"]:
            return True
            
        return False

    def _adjust_plate_color(self, plate_text, detected_color):
        """根据车牌长度调整颜色识别结果"""
        # 绿牌应该是8位，蓝牌是7位
        if detected_color == '绿牌' and len(plate_text) == 7 and plate_text[-1] not in ["学", "警"]:
            logger.info(f"根据长度将绿牌调整为蓝牌: {plate_text}")
            return "蓝牌"
        
        if detected_color != '绿牌' and len(plate_text) == 8 and plate_text[-1] not in ["学", "警"]:
            logger.info(f"根据长度将蓝牌调整为绿牌: {plate_text}")
            return "绿牌"
            
        return detected_color

    def _build_recognition_result(self, plate_text, plate_color, bbox, confidence, char_confidences):
        """构建识别结果字典"""
        # 字符校正
        corrected_text = plate_text.replace("O", '0').replace("I", '1')
        
        return {
            'plate': corrected_text,
            'color': plate_color,
            'box': bbox,
            'conf': confidence,
            'char_conf': char_confidences
        }

    def _passes_quality_check(self, char_confidences, min_confidence=60, max_low_conf_chars=1):
        """检查字符置信度是否满足质量要求"""
        low_confidence_count = sum(1 for conf in char_confidences if int(conf) < min_confidence)
        return low_confidence_count <= max_low_conf_chars

    def _visualize_results(self, image, results):
        """在图像上可视化识别结果"""
        for result in results:
            bbox = result['box']
            x1, y1, x2, y2 = [int(coord) for coord in bbox]
            
            # 绘制边界框
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
            
            # 添加车牌文本
            image = cv2ImgAddText(image, result['plate'], x1, y1-30, (255, 255, 0), 30)

    def _display_result_image(self, image):
        """显示识别结果图像"""
        try:
            resized_image = cv2.resize(image, (0, 0), fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
            cv2.imshow('车牌识别结果', resized_image)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        except Exception as e:
            logger.warning(f"显示图像失败: {e}")

# 保持向后兼容性的别名
LprConfig = LicensePlateConfig
lprservice = LicensePlateRecognitionService
scale_coords_landmarks = scale_landmark_coordinates
four_point_transform = apply_perspective_transform
order_points = sort_corner_points
opt = DEFAULT_CONFIG
